import os
import numpy as np
import torch
from torch import nn
import smplx
import open3d as o3d


class Simple_SMPLX(nn.Module):
    def __init__(self, smplx_model, init_smplx_params):
        super(Simple_SMPLX, self).__init__()
        self.smplx_model = smplx_model
        self.smplx_betas = nn.Parameter(init_smplx_params["betas"].clone().detach(), requires_grad=False)
        self.smplx_expression = nn.Parameter(init_smplx_params["expression"].clone().detach(), requires_grad=False)
        self.smplx_global_orient = nn.Parameter(init_smplx_params["global_orient"].clone().detach(), requires_grad=True)
        self.smplx_body_pose = nn.Parameter(init_smplx_params["body_pose"].clone().detach(), requires_grad=True)
        self.smplx_transl = nn.Parameter(init_smplx_params["transl"].clone().detach(), requires_grad=True)
    
    def forward(self):
        result_model = self.smplx_model(betas=self.smplx_betas, expression=self.smplx_expression, global_orient=self.smplx_global_orient, transl=self.smplx_transl, body_pose=self.smplx_body_pose, return_verts=True)
        result_vertices = result_model.vertices
        result_joints = result_model.joints
        return result_vertices, result_joints


def test(model_folder):
    smplx_model = smplx.create(model_folder, model_type="smplx",
                         gender="neutral", use_face_contour=False,
                         num_betas=10,
                         num_expression_coeffs=10,
                         ext="npz")
    smplx_model.to("cuda:0")

    print(smplx_model)

    # target model
    betas = torch.randn([1, smplx_model.num_betas], dtype=torch.float32).to("cuda:0")
    expression = torch.randn([1, smplx_model.num_expression_coeffs], dtype=torch.float32).to("cuda:0")
    global_orient = torch.randn([1, 3], dtype=torch.float32).to("cuda:0") * 0.1
    transl = torch.randn([1, 3], dtype=torch.float32).to("cuda:0") * 0.1

    target_model = smplx_model(betas=betas, expression=expression, global_orient=global_orient, transl=transl, return_verts=True)
    target_vertices = target_model.vertices.clone().detach()
    target_joints = target_model.joints.clone().detach()

    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(target_vertices.detach().cpu().numpy().squeeze())
    mesh.triangles = o3d.utility.Vector3iVector(smplx_model.faces)
    mesh.compute_vertex_normals()
    mesh.paint_uniform_color([0.3, 0.3, 0.3])
    o3d.io.write_triangle_mesh("./target.stl", mesh)

    init_smplx_params = {
        "betas": betas,
        "expression": expression,
        "global_orient": torch.randn([1, 3], dtype=torch.float32).to("cuda:0") * 0.1,
        "transl": torch.randn([1, 3], dtype=torch.float32).to("cuda:0") * 0.1,
        "body_pose": torch.rand([1, smplx_model.NUM_BODY_JOINTS, 3]).to("cuda:0") * 0.4 - 0.2,
    }
    
    optim_model = Simple_SMPLX(smplx_model=smplx_model, init_smplx_params=init_smplx_params)
    optim_model.to("cuda:0")

    optimizer = torch.optim.Adam(optim_model.parameters(), lr=0.002)
    optim_model.train()
    for epoch in range(1000):
        optimizer.zero_grad()
        v, j = optim_model()
        loss = torch.mean(torch.norm(v - target_vertices, dim=-1))
        # loss = torch.mean(torch.norm(j - target_joints, dim=-1))
        print(epoch, loss.item())
        loss.backward()
        optimizer.step()

        if epoch % 100 == 0:
            mesh = o3d.geometry.TriangleMesh()
            mesh.vertices = o3d.utility.Vector3dVector(v.detach().cpu().numpy().squeeze())
            mesh.triangles = o3d.utility.Vector3iVector(smplx_model.faces)
            mesh.compute_vertex_normals()
            mesh.paint_uniform_color([0.3, 0.3, 0.3])
            o3d.io.write_triangle_mesh("./epoch_{}.stl".format(epoch), mesh)


if __name__ == "__main__":
    # test
    model_folder = "/share/human_model/models"

    test(model_folder)
